import astro.continuum
import astro.spec
contpoints = [(5295.8048379364936, 48.324034395273664),
              (5338.9860724766522, 44.582893862361793),
              (5374.547089156782, 31.488901997170252),
              (5415.188251076931, 27.124238042106406),
              (5488.8503570572002, 20.889003820586623),
              (5850.9396763935338, 14.24525425351899),
              (6095.3834119637058, 12.480467960752287),
              (6314.2187561884302, 10.715681667985585),
              (6481.837317722262, 11.303943765574486),
              (4549.5677889294784, 44.540752279347387),
              (5159.5131100664776, 51.305766401619749),
              (5192.1056081425004, 73.659726109997976),
              (5243.322390833393, 70.130153524464575),
              (4759.0909908467675, 30.716592986008216),
              (4631.0490341195355, 37.187476059486123),
              (5546.5307055421827, 19.232772175064369),
              (5610.4012110915173, 20.165354634871093),
              (5686.0289738812771, 16.4350247956442),
              (5639.6354723379791, 16.4350247956442),
              (5584.0268232278622, 18.766480945161007),
              (5661.8789319820262, 16.4350247956442),
              (5985.1815637839718, 14.627656822897901),
              (5792.6714618920741, 15.444645388970331),
              (5742.2718214826073, 15.853139672006545),
              (4431.0408688578555, 74.289046002527044),
              (4354.0764919982703, 70.957310116831337),
              (4201.2166879577026, 70.957310116831337),
              (5276.5085534979589, 52.727081125733505),
              (5212.2114872067104, 98.759421403054702),
              (5052.0533402630572, 33.484053632754971),
              (5112.2587750630446, 40.653024659550894),
              (4982.4956049116172, 31.974796574482152),
              (4789.0198872534074, 30.842853780777542),
              (4863.2537728805755, 28.578968193368301),
              (4919.3675761893019, 31.220168045345748)]


class cofit(object):
    help_message = """
    Left click : add a new continuum point
    'd'        : delete the nearest point
    'q'        : quit
    """

    def makefig(self,spec, nbin=1):
        """ Set up the figure and do initial plots.

        Parameters
        ----------
        spec : dictionary with keys wa (wavelength), fl (flux)
               and er (1 sigma error).

        Returns
        -------
        f : the active figure.
        """
        wa,fl,er = [spec[k][0:-1:nbin] for k in 'wa fl er'.split()]
        f = pl.figure()
        f.clf()
        # axis for spectrum & continuum
        a0 = f.add_axes((0.05,0.1,0.9,0.6))
        a0.set_autoscale_on(0)
        # axis for residuals
        a1 = f.add_axes((0.05,0.75,0.9,0.2),sharex=a0)
        a1.set_autoscale_on(0)
        a1.axhline(0,color='g')
        a1.axhline(1,color='g',alpha=0.5)
        a1.axhline(-1,color='g',alpha=0.5)
        a1.axhline(2,color='g',linestyle='dotted',alpha=0.5)
        a1.axhline(-2,color='g',linestyle='dotted',alpha=0.5)
        m0, = a1.plot([0],[0],'ok',ms=3,alpha=0.3)
        a1.set_ylim(-4,4)
        a0.axhline(0, color='0.7')
        a0.plot(wa,er, color='crimson', alpha=0.8)
        a0.plot(wa,fl, 'b', linestyle='steps-mid')
        m1, = a0.plot([0],[0],'r',alpha=0.7)
        m2, = a0.plot([0],[0],'rx',ms=10,mew=1,picker=5)
        a0.set_xlim(min(wa),max(wa))
        ymin = -np.median(er)*5
        ymax = sorted(fl)[int(len(fl)*0.95)]*1.5
        a0.set_ylim(ymin, ymax)
        f.canvas.draw()
        markers = dict(contpoints=m2, cont=m1, resid=m0)
        return f, markers

    def updatefig(self,f):
        """ Calculates the new continuum, residuals and updates the plots.

        Needs external variables spec, markers and nbin."""
        global contpoints
        x,y = zip(*contpoints)
        spline = astro.continuum.InterpCubicSpline(x,y)
        cont = spline(spec['wa']) 
        resid = (spec['fl'] - cont) / spec['er']
        markers['contpoints'].set_data(zip(*contpoints))
        markers['cont'].set_data(spec['wa'][0:-1:nbin], cont[0:-1:nbin])
        markers['resid'].set_data(spec['wa'][0:-1:nbin], resid[0:-1:nbin])
        f.canvas.draw()

    def on_click(self,event):
        """ Add a new continuum point."""
        if f.canvas.toolbar._active in ('ZOOM','PAN'):  return
        if event.inaxes != f.axes[0]:  return
        # add a point to contpoints
        x,y = event.xdata,event.ydata
        if x not in zip(*contpoints)[0]:
            self.contpoints.append((x,y))
            updatefig(f)

    def on_keypress(self,event):
        """ Remove a continuum point."""
        global contpoints
        if f.canvas.toolbar._active in ('ZOOM','PAN'):
            print 'Must not be in zoom/pan mode'
            return
        if event.inaxes != f.axes[0]:  return
        if event.key == 'd':
            # remove a point from contpoints
            contx,conty = zip(*contpoints)
            sep = np.hypot(event.xdata - contx, event.ydata - conty)
            contpoints.remove((contpoints[sep.argmin()]))
            updatefig(f)
        elif event.key == '?':
            print help_message 
        elif event.key == 'q':
            disconnect(f)

    def disconnect(self,f):
        print 'Disconnecting...'
        f.canvas.mpl_disconnect(on_click)
        f.canvas.mpl_disconnect(on_keypress)

    def modifypoints(self,f):
        """add/remove continuum points."""
        print help_message
        f.canvas.mpl_connect('key_press_event',on_keypress)
        f.canvas.mpl_connect('button_press_event',on_click)

    if 1:
        #setup
        s = astro.spec.read('comb.fits')
        spec = dict(wa=s.wa,fl=s.fl,er=s.er)
        nbin = 8
        f,markers = makefig(spec,nbin=nbin)
        updatefig(f)
        outfile = modifypoints(f)    
        show()
